package com.duyus.ai.controller;

import com.duyus.ai.entity.vo.Result;
import com.duyus.ai.rag.MyKeywordEnricher;
import com.duyus.ai.rag.MyTokenTextSplitter;
import com.duyus.ai.repository.ChatHistoryRepository;
import com.duyus.ai.repository.MultiFileRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@Slf4j
@RequiredArgsConstructor
@RestController
@RequestMapping("/ai/pdf")
public class PdfController {

//    原本的 内存 json保存
//    private final FileRepository fileRepository;
    private final MultiFileRepository fileRepository;

//    private final VectorStore vectorStore;

    private final ChatClient pdfChatClient;

    private final ChatHistoryRepository chatHistoryRepository;

    private final VectorStore pgVectorVectorStore;

    private final MyTokenTextSplitter myTokenTextSplitter;

    private final MyKeywordEnricher myKeywordEnricher;

//    private final Advisor RagCloudAdvisor;

    @RequestMapping(value = "/chat", produces = "text/html;charset=utf-8")
    public Flux<String> chat(String prompt, String chatId) {
        // 1.找到会话文件
        Resource file = fileRepository.getFile(chatId);
        if (!file.exists()) {
            // 文件不存在，不回答
            throw new RuntimeException("会话文件不存在！");
        }
        // 2.保存会话id 用于后续识别当前会话属于哪个功能模块或文档
        chatHistoryRepository.save("pdf", chatId);

//        String context = pgVectorVectorStore.similaritySearch(SearchRequest.builder()
//                        .query(prompt)
//                        .topK(3)
//                        .build())
//                .stream()
//                .map(Document::getText)
//                .reduce("", (a, b) -> a + b + "\n");

        // 3.请求模型 构建并发起一个带有参数和过滤条件的流式聊天请求
        return pdfChatClient.prompt()
//                .system(context)
                .user(prompt)
                // 添加对话策略建议器（advisors）：设置当前会话 ID
                // CHAT_MEMORY_CONVERSATION_ID_KEY 是用于在聊天过程中识别会话上下文的键
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                // 添加另一个建议器：设置问题过滤表达式
                // 表示只从指定文件名（file.getFilename()）的文档中查找答案
//                .advisors(a -> a.param(FILTER_EXPRESSION, "file_name == '" + file.getFilename() + "'"))
                // 应用 RAG 检索增强服务（基于云知识库服务）
//                .advisors(RagCloudAdvisor)
                // 应用 RAG 检索增强服务（基于 PgVector 向量存储）
//                .advisors(new QuestionAnswerAdvisor(pgVectorVectorStore))
                // 应用自定义的 RAG 检索增强服务（文档查询器 + 上下文增强器）
//                .advisors(
//                        RagCustomAdvisorFactory.AppRagCustomAdvisor(
//                                pgVectorVectorStore, "单身"
//                        )
//                )
                .advisors(new QuestionAnswerAdvisor(
                        pgVectorVectorStore,
                        SearchRequest.builder()
                                .query(prompt) // 当前用户的问题
                                .similarityThreshold(0.3)
                                .topK(3)
                                .build()
                ))
                .stream()
                .content();
    }

    /**
     * 文件上传
     */
    @RequestMapping("/upload/{chatId}")
    public Result uploadPdf(@PathVariable String chatId, @RequestParam("file") MultipartFile file) {
        try {
            // 1. 校验文件是否为PDF格式
            if (!Objects.equals(file.getContentType(), "application/pdf")) {
                return Result.fail("只能上传PDF文件！");
            }
            // 2.保存文件
            boolean success = fileRepository.save(chatId, file.getResource());
            if (!success) {
                return Result.fail("保存文件失败！");
            }
            // 3.写入向量库
            this.writeToVectorStore(file.getResource());
            return Result.ok();
        } catch (Exception e) {
            log.error("Failed to upload PDF.", e);
            return Result.fail("上传文件失败！");
        }
    }

    /**
     * 文件下载
     */
    @GetMapping("/file/{chatId}")
    public ResponseEntity<Resource> download(@PathVariable("chatId") String chatId) throws IOException {
        // 1.读取文件
        Resource resource = fileRepository.getFile(chatId);
        if (!resource.exists()) {
            return ResponseEntity.notFound().build();
        }
        // 2.文件名编码，写入响应头
        String filename = URLEncoder.encode(Objects.requireNonNull(resource.getFilename()), StandardCharsets.UTF_8);
        // 3.返回文件
        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_OCTET_STREAM)
                .header("Content-Disposition", "attachment; filename=\"" + filename + "\"")
                .body(resource);
    }

    private void writeToVectorStore(Resource resource) {
        // 1.创建PDF的读取器
//        ParagraphPdfDocumentReader reader = new ParagraphPdfDocumentReader(resource,
//                PdfDocumentReaderConfig.builder() // 每一个标题作为一个Document
//                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
//                        .build()
//        );
        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                resource, // 文件源
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1) // 每1页PDF作为一个Document
                        .build()
        );
        // 2.读取PDF文档，拆分为Document
        List<Document> documents = reader.read();
        // 3.自主切分文档
        List<Document> splitDocuments = myTokenTextSplitter.splitDocuments(documents);
        // 4.自动补充关键词元信息
        List<Document> enrichedDocuments = myKeywordEnricher.enrichDocuments(splitDocuments);
        // 3.写入向量库
        pgVectorVectorStore.add(enrichedDocuments);
    }
}